import os
import argparse
from subprocess import Popen, PIPE
from uuid import uuid1
import sys
import pandas as pd
import scipy.stats
import numpy as np
from sklearn.metrics import mean_squared_error
import math
import gzip

sep = "||"
seeds = [111, 222, 333, 444, 555, 666, 777, 888, 999, 1111,
         2222, 3333, 4444, 5555, 6666, 7777, 8888, 9999, 11111]


def run_cmd(args_list):
    proc = Popen(args_list, shell=True, stdout=PIPE, stderr=PIPE)
    stdinfo = proc.communicate()
    # print(stdinfo)
    return stdinfo, proc.returncode


def exe_one_cmd(cmd_str, namestr):
    sys.stderr.write("{} cmd: {}\n".format(namestr, cmd_str))
    sys.stderr.flush()
    stdinfo, returncode = run_cmd(cmd_str)
    stdout, stderr = stdinfo
    sys.stderr.write("stdout:\n{}\n,stderr:\n{}\n".format(stdout, stderr))
    return returncode


def exe_mod2freq_cmd(mod_file, freq_file, prob_cf):
    mod2freq_cmd = "python ~/methylation/deepsignal3/scripts/call_modification_frequency.py --input_path {} --result_file {} --bed --sort --prob_cf {}".format(mod_file, freq_file, prob_cf)
    returncode = exe_one_cmd(mod2freq_cmd, "mod2freq_cmd")
    return returncode


def exe_mod2freq_cmd_deepmod2(mod_file, genome, motif):
    basedir = os.path.dirname(mod_file)
    fname, _ = os.path.splitext(os.path.basename(mod_file))
    freq_file = basedir +"/" + fname + ".per_site"
    mod2freq_cmd = "python ~/tools/DeepMod2/deepmod2 merge --input {} --output {} --prefix {}".format(mod_file, basedir, fname)
    returncode = exe_one_cmd(mod2freq_cmd, "mod2freq_cmd_dm2 step1")
    if returncode:
        raise ValueError("mod2freq_cmd wrong!")
    
    freq2bed_cmd = "python ~/tools/ont_5mC_2/scripts/persite2bedmethyl_deepmod2.py --per_site {} --genome {} --motif {}".format(freq_file, genome, motif)
    returncode = exe_one_cmd(freq2bed_cmd, "mod2freq_cmd_dm2 step2")
    freq_file2 = freq_file + ".bed"
    os.remove(freq_file)
    return returncode, freq_file2


def exe_mod2freq_cmd_rockfish(mod_file, freq_file, prob_cf):
    mod2freq_cmd = "python ~/tools/ont_5mC_2/scripts/call_modification_frequency_rf.py --input_path {} --result_file {} --bed --sort --prob_cf {}".format(mod_file, freq_file, prob_cf)
    returncode = exe_one_cmd(mod2freq_cmd, "mod2freq_cmd_rf")
    return returncode


def exe_mod2freq_cmd_deepplant(mod_file, freq_file, methyl_th_dp, prob_cf):
    mod2freq_cmd = "python ~/tools/ont_5mC_2/scripts/call_modification_frequency_dp.py --input_path {} --result_file {} --bed --sort --prob_cf {} --methyl_threshold {}".format(mod_file, freq_file, prob_cf, methyl_th_dp)
    returncode = exe_one_cmd(mod2freq_cmd, "mod2freq_cmd_dp")
    if returncode:
        raise ValueError("mod2freq_cmd wrong!")
    return returncode


def gen_freq_bed(mods, fa, genome_size, cov, seedtmp, rid_col=5, prob_cf=0.0, 
                 deepmod2=False, rockfish=False, deepplant=False, 
                 genome=None, motif=None, methyl_th_dp=0.5):
    uuidstr = str(uuid1())
    fname, fext = os.path.splitext(fa)
    sub_fafile = fname + ".samp_cov_{}.{}.fa".format(cov, uuidstr)
    sampfa_cmd = "rasusa reads --coverage {} --genome-size {} " \
                 "-o {} --seed {} {}".format(cov, genome_size, sub_fafile, 
                                             seedtmp, fa)
    returncode = exe_one_cmd(sampfa_cmd, "sampfa_cmd")
    if returncode:
        raise ValueError("sampfa_cmd wrong!")
    
    sub_fa_rid_file = sub_fafile + ".rid"
    extract_rid_cmd = "seqkit seq -n {} > {}".format(sub_fafile, sub_fa_rid_file)
    returncode = exe_one_cmd(extract_rid_cmd, "extract_rid_cmd")
    if returncode:
        raise ValueError("extract_rid_cmd wrong!")
    
    fname, fext = os.path.splitext(mods)
    samp_mod_file = fname + ".samp_cov_{}.{}.tsv".format(cov, uuidstr)
    extract_mod_cmd = "awk 'BEGIN{{ FS=OFS=\"\t\" }} NR==FNR {{a[$1]; next}} ${} in a {{print $0}}' {} {} > {}".format(
        rid_col, sub_fa_rid_file, mods, samp_mod_file)
    returncode = exe_one_cmd(extract_mod_cmd, "extract_mod_cmd")
    if returncode:
        raise ValueError("extract_mod_cmd wrong!")
    
    if deepmod2:
        returncode, freq_bed = exe_mod2freq_cmd_deepmod2(samp_mod_file, genome, motif)
        if returncode:
            raise ValueError("mod2freq_cmd wrong! in deepmod2")
    elif rockfish:
        fname, fext = os.path.splitext(samp_mod_file)
        freq_bed = fname + ".freq.bed"
        returncode = exe_mod2freq_cmd_rockfish(samp_mod_file, freq_bed, prob_cf)
        if returncode:
            raise ValueError("mod2freq_cmd wrong! in rockfish")
    elif deepplant:
        fname, fext = os.path.splitext(samp_mod_file)
        freq_bed = fname + ".th{}.freq.bed".format(methyl_th_dp)
        returncode = exe_mod2freq_cmd_deepplant(samp_mod_file, freq_bed, methyl_th_dp, prob_cf)
        if returncode:
            raise ValueError("mod2freq_cmd wrong! in deepplant")
    else:
        fname, fext = os.path.splitext(samp_mod_file)
        freq_bed = fname + ".freq.bed"
        returncode = exe_mod2freq_cmd(samp_mod_file, freq_bed, prob_cf)
        if returncode:
            raise ValueError("mod2freq_cmd wrong!")
    
    os.remove(sub_fafile)
    os.remove(sub_fa_rid_file)
    os.remove(samp_mod_file)
    return freq_bed


def read_methylbed2(freqfile, contig_prefix, contig_names, cov_cf):
    # methylbed format
    # "chromosome", "pos", "end", "na1", "na2", "strand", "na3", "na4", "na5", "coverage", "rpercent"
    contigset = set(contig_names.strip().split(",")) if contig_names is not None else None
    freqinfo = {}
    covs = []
    if freqfile.endswith(".gz"):
        infile = gzip.open(freqfile, 'rt')
    else:
        infile = open(freqfile, 'r')
    for line in infile:
        words = line.strip().split("\t")
        chrom = words[0]
        # m_key = "\t".join([words[0], words[1], words[5]])
        m_key = "\t".join([words[0], words[1]])  # discard strand
        cov = float(words[9])
        rmet = float(words[10]) / 100
        # methy_cov = rmet * cov

        cnt_flag = 0
        if contig_prefix is not None:
            if str(chrom).startswith(contig_prefix):
                cnt_flag = 1
        elif contig_names is not None:
            if chrom in contigset:
                cnt_flag = 1
        else:
            cnt_flag = 1

        if cnt_flag == 1:
            covs.append(cov)
            if cov >= cov_cf:
                freqinfo[m_key] = rmet
    infile.close()
    return np.mean(covs) if len(covs) > 0 else 0, freqinfo


def cal_corr_df1_vs_df2_v2(freqinfo1, freqinfo2):
    keys_inter = set(freqinfo1.keys()).intersection(set(freqinfo2.keys()))
    keys_inter = sorted(list(keys_inter))
    rmet1, rmet2 = [], []
    for ktmp in keys_inter:
        rmet1.append(freqinfo1[ktmp])
        rmet2.append(freqinfo2[ktmp])
    rmet1 = np.array(rmet1)
    rmet2 = np.array(rmet2)
    if len(rmet1) > 1 and len(rmet2) > 1:
        corr, pvalue = scipy.stats.pearsonr(rmet1, rmet2)
        scorr, _ = scipy.stats.spearmanr(rmet1, rmet2)  # spearman
        _, _, r_value, _, _ = scipy.stats.linregress(rmet1, rmet2)
        r_square = r_value ** 2  # coefficient of determination
        rmse = math.sqrt(mean_squared_error(rmet1, rmet2))  # RMSE
    else:
        corr, r_square, scorr, rmse = 0, 0, 0, 0

    return len(freqinfo1.keys()), len(freqinfo2.keys()), len(keys_inter), \
        corr, scorr, r_square, rmse


def display_args(args, is_stderr=False):
    def print_outputstr(outstr):
        if is_stderr:
            sys.stderr.write(outstr + "\n")
        else:
            sys.stdout.write(outstr + "\n")
    arg_vars = vars(args)
    outputstr = "# ===============================================\n## parameters: "
    print_outputstr(outputstr)
    for arg_key in arg_vars.keys():
        if arg_key != 'func':
            outputstr = "{}:\n\t{}".format(arg_key, arg_vars[arg_key])
            print_outputstr(outputstr)
    print_outputstr("# ===============================================")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mods", type=str, required=True, help="per-read call_mods file")
    parser.add_argument("--fa", type=str, required=True, help="fasta file")
    parser.add_argument("--covs", type=str, required=False, default="10,20,30,40,50",
                        help="")
    parser.add_argument("--genome_size", type=str, required=False, default="251.49MB",
                        help="default 251.49MB (human chr1)")
    parser.add_argument("--repeat", type=int, required=False, default=5, help="repeat times for a coverage, "
                                                                              "default 5")
    parser.add_argument("--total", action="store_true", default=False, help="if eval all reads in bam")

    parser.add_argument("--cmp_bed", action="append", required=True, help="cmp bed")
    # parser.add_argument("--contig_names", type=str, required=False, default="chr1,chr2,chr3,chr4,chr5,chr6,chr7,"
    #                                                                         "chr8,chr9,chr10,chr11,chr12,chr13,"
    #                                                                         "chr14,chr15,chr16,chr17,chr18,chr19,"
    #                                                                         "chr20,chr21,chr22,chrX,chrY")
    parser.add_argument("--contig_names", type=str, required=False, default="chr1")
    parser.add_argument("--cov_cf", type=int, required=False, default=5, help="default 5")

    parser.add_argument("--rid_col", type=int, required=False, default=5, help="default 5")
    parser.add_argument('--prob_cf', type=float, action="store", required=False, default=0.0,
                        help='this is to remove ambiguous calls (only for count-mode now). '
                             'if abs(prob1-prob0)>=prob_cf, then we use the call. e.g., proc_cf=0 '
                             'means use all calls. range [0, 1], default 0.0.')
    parser.add_argument("--deepmod2", action="store_true", default=False, help="if deepmod2 format")
    parser.add_argument("--rockfish", action="store_true", default=False, help="if rockfish format")
    parser.add_argument("--deepplant", action="store_true", default=False, help="if deepplant format")

    parser.add_argument("--genome", type=str, required=False, help="genome reference for deepmod2")
    parser.add_argument("--motif", type=str, required=False, default="CG", help="motif for deepmod2")

    parser.add_argument("--methyl_th", type=float, required=False, default=0.5,
                        help="methylation threshold for deepplant, default 0.5")

    args = parser.parse_args()
    assert args.deepmod2 & args.rockfish == 0, "deepmod2 and rockfish can't be both True"
    assert args.deepmod2 & args.deepplant == 0, "deepmod2 and deepplant can't be both True"
    assert args.rockfish & args.deepplant == 0, "rockfish and deepplant can't be both True"

    display_args(args)

    sys.stdout.write("eval modsfile - {} ====\n".format(args.mods))
    sys.stdout.write("fa file to shuffle - {} ====\n".format(args.fa))
    
    sys.stdout.write("\ncmp files:\n")
    cmpfname2rmetinfo = dict()
    for cmp_file in args.cmp_bed:
        cmpmean_cov, cmpfname2rmetinfo[os.path.basename(cmp_file)] = read_methylbed2(cmp_file,
                                                                                     None,
                                                                                     args.contig_names,
                                                                                     args.cov_cf)
        sys.stdout.write("cmpfile: {}, mean_covarge: {}\n".format(cmp_file, cmpmean_cov))
    sys.stdout.flush()

    covs = [int(x) for x in args.covs.strip().split(",")]
    for cov in covs:
        sys.stdout.write("\nfor cov-{} ====\n".format(cov))
        freq_beds = []
        for ridx in range(args.repeat):
            seedtmp = seeds[ridx]
            freq_bed = gen_freq_bed(args.mods, args.fa, args.genome_size, cov, seedtmp, 
                                    args.rid_col, args.prob_cf, args.deepmod2, args.rockfish, 
                                    args.deepplant, args.genome, args.motif, args.methyl_th)
            freq_beds.append(freq_bed)
        
        c_covs, c_rmets = [], []
        for cbedtmp in freq_beds:
            mean_cov, rmetinfo = read_methylbed2(cbedtmp, None, args.contig_names, args.cov_cf)
            c_covs.append(mean_cov)
            c_rmets.append(rmetinfo)
        sys.stdout.write("means: {}\n".format(c_covs)) 

        for cmp_file in cmpfname2rmetinfo.keys():
            sys.stdout.write("== cmp_file: {}\n".format(cmp_file))
            sys.stdout.write("\t".join(["cmp_file", "cmpnum", "tgsnum", "internum",
                                        "pearson", "rsquare", "spearman", "RMSE"]) + "\n")
            cmp_df = cmpfname2rmetinfo[cmp_file]
            sitenums_inter = []
            sitenums_pb = []
            sitenums_cmp = []
            corrs_pearson = []
            corrs_spearman = []
            corrs_rsquare = []
            corrs_rmse = []
            for c_rmet_tmp in c_rmets:
                smrtnum, cmpnum, internum, \
                    pcorr, scorr, r_square, rmse = cal_corr_df1_vs_df2_v2(c_rmet_tmp, cmp_df)
                sitenums_inter.append(internum)
                sitenums_pb.append(smrtnum)
                sitenums_cmp.append(cmpnum)
                corrs_pearson.append(pcorr)
                corrs_spearman.append(scorr)
                corrs_rsquare.append(r_square)
                corrs_rmse.append(rmse)
                sys.stdout.write("{}\t{}\t{}\t{}\t{:.4f}\t{:.4f}\t{:.4f}\t{:.4f}\n".format(cmp_file, cmpnum, smrtnum,
                                                                              internum, pcorr, r_square,
                                                                              scorr, rmse))
            sys.stdout.write("{}\t{:.2f}\t{:.2f}\t{:.2f}\t{:.4f}\t{:.4f}\t{:.4f}\t{:.4f}\n".format("mean",
                                                                                      np.mean(sitenums_cmp),
                                                                                      np.mean(sitenums_pb),
                                                                                      np.mean(sitenums_inter),
                                                                                      np.mean(corrs_pearson),
                                                                                      np.mean(corrs_rsquare),
                                                                                      np.mean(corrs_spearman),
                                                                                      np.mean(corrs_rmse)
                                                                                      ))
            sys.stdout.write("{}\t{:.2f}\t{:.2f}\t{:.2f}\t{:.4f}\t{:.4f}\t{:.4f}\t{:.4f}\n".format("SD",
                                                                                                   np.std(
                                                                                                       sitenums_cmp),
                                                                                                   np.std(sitenums_pb),
                                                                                                   np.std(
                                                                                                       sitenums_inter),
                                                                                                   np.std(
                                                                                                       corrs_pearson),
                                                                                                   np.std(
                                                                                                       corrs_rsquare),
                                                                                                   np.std(
                                                                                                       corrs_spearman),
                                                                                                   np.std(corrs_rmse)
                                                                                                   ))
        sys.stdout.flush()
        del c_rmets
    
    if args.total:
        sys.stdout.write("\nfor total reads ====\n")
        if args.deepmod2:
            returncode, freq_bed_total = exe_mod2freq_cmd_deepmod2(args.mods, args.genome, args.motif)
            if returncode:
                raise ValueError("mod2freq_cmd wrong! in deepmod2")
        elif args.rockfish:
            fname, fext = os.path.splitext(args.mods)
            freq_bed_total = fname + ".freq.bed"
            returncode = exe_mod2freq_cmd_rockfish(args.mods, freq_bed_total, args.prob_cf)
            if returncode:
                raise ValueError("mod2freq_cmd wrong! in rockfish")
        elif args.deepplant:
            fname, fext = os.path.splitext(args.mods)
            freq_bed_total = fname + ".th{}.freq.bed".format(args.methyl_th)
            returncode = exe_mod2freq_cmd_deepplant(args.mods, freq_bed_total, args.methyl_th, args.prob_cf)
            if returncode:
                raise ValueError("mod2freq_cmd wrong! in deepplant")
        else:
            fname, fext = os.path.splitext(args.mods)
            freq_bed_total = fname + ".freq.bed"
            returncode = exe_mod2freq_cmd(args.mods, freq_bed_total, args.prob_cf)
            if returncode:
                raise ValueError("mod2freq_cmd wrong! in total")
        mean_cov, rmetinfo = read_methylbed2(freq_bed_total, None, args.contig_names, args.cov_cf)
        sys.stdout.write("means: {}\n".format(mean_cov))
        for cmp_file in cmpfname2rmetinfo.keys():
            sys.stdout.write("== cmp_file: {}\n".format(cmp_file))
            sys.stdout.write("\t".join(["cmp_file", "cmpnum", "tgsnum", "internum",
                                        "pearson", "rsquare", "spearman", "RMSE"]) + "\n")
            cmp_df = cmpfname2rmetinfo[cmp_file]
            smrtnum, cmpnum, internum, \
                pcorr, scorr, r_square, rmse = cal_corr_df1_vs_df2_v2(rmetinfo, cmp_df)
            sys.stdout.write("{}\t{}\t{}\t{}\t{:.4f}\t{:.4f}\t{:.4f}\t{:.4f}\n".format(cmp_file, cmpnum, smrtnum,
                                                                                       internum, pcorr, r_square,
                                                                                       scorr, rmse))


if __name__ == "__main__":
    main()
